import torch
import numpy as np
from scipy.interpolate import CubicSpline

class Gen3DTrajectories:
    def __init__(self, num_objects=128, num_points=10, device='cuda'):
        self.num_objects = num_objects
        self.num_points = num_points
        self.device = device
        
        # 三维运动参数
        self.speed = torch.zeros(num_objects, device=device)
        self.direction = torch.ones(num_objects, device=device)
        self.current_t = torch.zeros(num_objects, device=device)
        
        # 生成三维轨迹
        self.trajectories, self.tangents = self.gen_trajectories()
    
    def gen_trajectories(self):
        """生成三维样条轨迹"""
        control_points = torch.rand(self.num_objects, self.num_points, 3, device=self.device) * 15 + 5.0
        
        traj_list, tangent_list = [], []
        for i in range(self.num_objects):
            np_points = control_points[i].cpu().numpy()
            t_param = np.linspace(0, 1, self.num_points)
            
            # 三轴独立样条
            cs_x = CubicSpline(t_param, np_points[:, 0], bc_type='natural')
            cs_y = CubicSpline(t_param, np_points[:, 1], bc_type='natural')
            cs_z = CubicSpline(t_param, np_points[:, 2], bc_type='natural')
            
            t_dense = np.linspace(0, 1, 1000)
            traj = torch.tensor(np.column_stack((cs_x(t_dense), cs_y(t_dense), cs_z(t_dense))),
                               dtype=torch.float32, device=self.device)
            
            # 三维切线方向
            dx = cs_x(t_dense, 1)
            dy = cs_y(t_dense, 1)
            dz = cs_z(t_dense, 1)
            tangent = torch.tensor(np.column_stack((dx, dy, dz)),
                                  dtype=torch.float32, device=self.device)
            
            traj_list.append(traj)
            tangent_list.append(tangent)
            
        return torch.stack(traj_list), torch.stack(tangent_list)
    
    def update(self):
        """更新三维位姿"""
        self.speed = torch.rand(self.num_objects, device=self.device) * 0.001
        self.current_t += self.speed * self.direction
        
        # 边界处理（保持与二维相同）
        overflow_mask = (self.current_t >= 1.0) | (self.current_t <= 0.0)
        self.direction[overflow_mask] *= -1
        self.current_t = torch.clamp(self.current_t, 0.0, 1.0)
        
        indices = (self.current_t * (self.trajectories.shape[1] - 1)).long()
        positions = self.trajectories[torch.arange(self.num_objects, device=self.device), indices]
        
        # 三维朝向计算
        tangents = self.tangents[torch.arange(self.num_objects, device=self.device), indices]
        adjusted_tangents = tangents * self.direction.unsqueeze(-1)
        
        # 归一化处理
        lengths = torch.norm(adjusted_tangents, dim=1, keepdim=True)
        mask = (lengths < 1e-6).squeeze()
        adjusted_tangents[mask] = torch.tensor([1.0, 0.0, 0.0], device=self.device)
        adjusted_tangents = adjusted_tangents / (lengths + 1e-6)
        
        # 调用三维四元数计算
        quats = self._compute_3d_quaternion(adjusted_tangents)
        return positions, quats
    
    def _compute_3d_quaternion(self, adjusted_tangents):
        # 构建前方向（切线）和上方向（世界Z轴）
        forward = torch.nn.functional.normalize(adjusted_tangents, dim=1)
        up = torch.tensor([[0.0, 0.0, 1.0]], 
                     device=self.device, 
                     dtype=torch.float32)

        # 计算右方向
        right = torch.cross(up, forward, dim=1)
        up = torch.cross(forward, right, dim=1)
        
        # 构建旋转矩阵
        rot_matrix = torch.stack((right, up, forward), dim=2)
        
        # 转换为四元数
        trace = rot_matrix[:,0,0] + rot_matrix[:,1,1] + rot_matrix[:,2,2]
        q = torch.zeros((self.num_objects, 4), device=self.device)
        
        q[:, 0] = 0.5 * torch.sqrt(1 + trace)
        q[:, 1] = (rot_matrix[:,2,1] - rot_matrix[:,1,2]) / (4 * q[:, 0])
        q[:, 2] = (rot_matrix[:,0,2] - rot_matrix[:,2,0]) / (4 * q[:, 0])
        q[:, 3] = (rot_matrix[:,1,0] - rot_matrix[:,0,1]) / (4 * q[:, 0])
        return q